#include <stdio.h>
#include <time.h>

#define abs(x) (x >= 0 ? x : -x)
#define max(x, y) (x > y ? x : y)

#define BLOCK_SIZE 16

typedef struct {
    int order;
    float * elements;
} Matrix;

__global__ void MVKernel_shm(Matrix A, Matrix X, Matrix Y)
{
    __shared__ float Xds[BLOCK_SIZE];

    int bx = blockIdx.x;
    //int by = blockIdx.y;
    int tx = threadIdx.x;
    //int ty = threadIdx.y;
    // Calculate the row index
    //int Row = by * BLOCK_SIZE + ty;
    // Calculate the column index
    int Row = bx * BLOCK_SIZE + tx;

    float Pvalue = 0;

    for (unsigned int m = 0; m < (A.width-1)/BLOCK_SIZE+1; ++m)
    {
      if(m*BLOCK_SIZE + tx < A.width)
        Xds[tx] = X.elements[m*BLOCK_SIZE + tx];
      else
        Xds[tx] = 0;
      __syncthreads();

      for (unsigned int k = 0; k < BLOCK_SIZE; k++)
             if(Row<A.height && m*BLOCK_SIZE +k<A.width)
                Pvalue += A.elements[m*BLOCK_SIZE+Row*A.width+k] * Xds[k];
        __syncthreads();
    }

    if(Row < A.height)
    Y.elements[Row] = Pvalue;
    __syncthreads();
}

void MatMul(const Matrix A, const Matrix B, Matrix C)
{
    Matrix d_A;
    d_A.order = A.order;
    size_t size = A.order * A.order * sizeof(float);
    cudaMalloc(&d_A.elements, size);
    cudaMemcpy(d_A.elements, A.elements, size, cudaMemcpyHostToDevice);

    Matrix d_B;
    d_B.order = B.order;
    size = B.order * A.order * sizeof(float);
    cudaMalloc(&d_B.elements, size);
    cudaMemcpy(d_B.elements, B.elements, size, cudaMemcpyHostToDevice);

    Matrix d_C;
    d_C.order = C.order;
    size = C.order * A.order * sizeof(float);
    cudaMalloc(&d_C.elements, size);

    dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 dimGrid((A.order + dimBlock.x - 1) / dimBlock.x,
            (A.order + dimBlock.y - 1) / dimBlock.y);

    MVKernel_gm<<<dimGrid, dimBlock>>>(d_A, d_B, d_C);
    cudaThreadSynchronize();

    cudaMemcpy(C.elements, d_C.elements, size, cudaMemcpyDeviceToHost);

    cudaFree(d_A.elements);
    cudaFree(d_B.elements);
    cudaFree(d_C.elements);
}


int main(int argc, char * argv[])
{
    FILE * ifp;
    Matrix A, B, C;
    int a1 = 3;
    int i, j;

    A.order = a1;
    A.elements = (float *) malloc(A.order * A.order * sizeof(float));
    B.order = 1;
    B.elements = (float *) malloc(A.order * B.order * sizeof(float));
    C.order = 1;
    C.elements = (float *) malloc(A.order * C.order * sizeof(float));

    for(i = 0; i < A.order; i++)
        for(j = 0; j < A.order; j++)
            scanf("%f", &A.elements[i * A.order + j]);

    for(i = 0; i < A.order; i++)
        scanf("%f", &B.elements[i]);

    MatMul(A, B, C);

    for(i = 0; i < A.order; i++)
    {
        for(j = 0; j < A.order; j++)
            printf("%f ", A.elements[i * A.order + j]);
        printf("\n");
    }

    for(i = 0; i < A.order; i++)
    {
        printf("%f\n", B.elements[i]);
    }

    for(i = 0; i < A.order; i++)
    {
        printf("%f\n", C.elements[i]);
    }
}
